import torch
from model.VGG16 import Vgg16
from model.ResNet import ResNet50,ResNet101,ResNet152
import torch.optim.lr_scheduler as lr_scheduler


def get_model(name,device):
    if name == 'Vgg16':
        model = Vgg16().to(device)
    elif name == 'ResNet50':
        model = ResNet50().to(device)
    elif name == 'ResNet101':
        model = ResNet101().to(device)
    elif name == 'ResNet152':
        model = ResNet152().to(device)
    else:
        model = None
    return model


def get_train_instance(arg,model,device):
    # 损失函数
    criterion = torch.nn.CrossEntropyLoss().to(device)
    # 优化器
    if arg['optim'] == 'Adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=float(arg['lr']))
    elif arg['optim'] == 'SGD':
        optimizer = torch.optim.SGD(model.parameters(), lr=float(arg['lr']), momentum=arg['momentum'],
                                    weight_decay=arg['weight_decay'])
    else:
        raise ValueError("请检查配置文件中optim是否正确，你的输入为{}，应为Adam or SGD".format(arg['optim']))
    # scheduler
    if arg['scheduler']:
        scheduler = lr_scheduler.StepLR(optimizer, step_size=arg['step_size'], gamma=arg['gamma'])
    else:
        scheduler = None
    return (criterion, optimizer, scheduler)

def train_step(loader,mo):
    pass